import abc
from collections.abc import Iterable
from typing import Any, Literal, TypeVar, Union

import numpy as np
import torch
from numpy.typing import ArrayLike
from zkl_aiutils_training import DockPlugin
from zkl_ptutils_metrics import Metric

AnyInput = TypeVar('AnyInput')
AnyOutput = TypeVar('AnyOutput')
MetricScope = Literal['epoch', 'run']
MetricValue = Union[torch.Tensor, ArrayLike, Any]
MetricOrValue = Union[Metric, MetricValue]
MetricScopeFilter = MetricScope | Iterable[MetricScope] | None


class MetricsMixin(abc.ABC):
    @abc.abstractmethod
    def _get_scoped_metrics(self) -> dict[str, tuple[MetricOrValue, MetricScope]]:
        pass

    def _get_metrics(self, scope: MetricScopeFilter = None) -> dict[str, MetricOrValue]:
        scopes = None if scope is None \
            else (scope,) if isinstance(scope, str) \
            else tuple(scope)
        metrics = {}
        for name, (metric, scope) in self._get_scoped_metrics().items():
            if scopes is None or scope in scopes:
                metrics[name] = metric
        return metrics

    def _get_metrics_value(self, scope: MetricScopeFilter = None) -> dict[str, MetricValue]:
        metrics_value = {}
        for name, metric in self._get_metrics(scope).items():
            metrics_value[name] = get_metric_value(metric)
        return metrics_value

    def _get_metrics_scalar_value(self, scope: MetricScopeFilter = None) -> dict[str, float | int]:
        metrics_scalar_value = {}
        for name, metric in self._get_metrics(scope).items():
            scalar_value = get_metric_scalar_value(metric)
            if scalar_value is not None:
                metrics_scalar_value[name] = scalar_value
        return metrics_scalar_value

    def _get_scoped_metric(self, name: str) -> tuple[MetricOrValue, MetricScope] | None:
        return self._get_scoped_metrics().get(name)

    def _get_metric(self, name: str) -> MetricOrValue | None:
        scoped_metric = self._get_scoped_metric(name)
        return scoped_metric[0] if scoped_metric is not None else None

    def _get_metric_value(self, name: str) -> MetricValue | None:
        metric = self._get_metric(name)
        return get_metric_value(metric)

    def _get_metric_scalar_value(self, name: str) -> float | int | None:
        value = self._get_metric_value(name)
        value = get_metric_scalar_value(value)
        return value


class MetricsProvider(MetricsMixin, DockPlugin, abc.ABC):
    def _ensure_after_collectors(self):
        self.dock.ensure_after_plugins(self.dock.iter_plugins(MetricsCollector))


class MetricsCollector(MetricsMixin, DockPlugin, abc.ABC):
    def _ensure_after_providers(self):
        self.dock.ensure_after_plugins(self.dock.iter_plugins(MetricsProvider))

    def _get_scoped_metrics(self) -> dict[str, tuple[Metric, MetricScope]]:
        return self.dock.scan_agg_plugins(
            self.dock.iter_plugins(MetricsProvider), {},
            lambda plugin, metrics: {**metrics, **plugin._get_scoped_metrics()})


def get_metric_value(metric: MetricOrValue | None) -> MetricValue:
    if isinstance(metric, Metric):
        return metric.compute()
    return metric


def get_metric_scalar_value(metric: MetricOrValue | None) -> float | int | None:
    if isinstance(metric, Metric):
        value = metric.compute()
        return get_metric_scalar_value(value)
    if isinstance(metric, torch.Tensor):
        value = metric.detach().cpu().numpy()
        return get_metric_scalar_value(value)
    if isinstance(metric, np.ndarray):
        if metric.shape == ():
            value = metric.item()
            return get_metric_scalar_value(value)
    if isinstance(metric, (int, float)):
        return metric
    return None
